#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<ctime>
#include<iostream>
#include<map>
#include<iomanip>
#include"Util.h"
#include"KBestParseForest.h"
using namespace std;

BestPair::BestPair(int first, int second){
	this->first = first;
	this->second = second;
}

BestPair::BestPair(){}

int KBestParseForest::rootType = 0;

void KBestParseForest::setRootType(int i){
	rootType = i;
}

KBestParseForest::KBestParseForest(){}

KBestParseForest::KBestParseForest(int start, int end, DependencyInstance &inst, int K){
	this->start = start;
	this->end = end;
	this->K = K;
	this->sent = inst.forms;
	this->pos = inst.postags;
	this->chart = new ParseForestItem[(end + 1) * (end + 1) * 2 * 2 * K];

	for(int i = 0; i< end + 1; i++){
		for(int j = 0; j < end + 1; j++){
			for(int k = 0; k < 2; k++){
				for(int l = 0; l < 2; l++){
					for(int m = 0; m < K; m++){
						int key = getKey(i, j, k, l, m);
						chart[key] = ParseForestItem(i, -1, j, -1, k, l, Negative_Infinity, NULL, NULL);
					}
				}
			}
		}
	}
}

int KBestParseForest::getKey(int s, int t, int dir, int comp, int k){
	int key = s;
	key = key * (end + 1) + t;
	key = key * 2 + dir;
	key = key * 2 + comp;
	key = key * K + k;
	return key;
}

bool KBestParseForest::add(int s, int type, int dir, double score){
	bool added = false;

	int key = getKey(s, s, dir, 1, K-1);
	if(chart[key].prob > score){
		return false;
	}
	for(int i = 0; i < K; i++){
		key = getKey(s, s, dir, 1, i);
		if(chart[key].prob < score){
			ParseForestItem tmp = chart[key];
			chart[key] = ParseForestItem(s, type, dir, score);
			for(int j = i + 1; j < K && tmp.prob > Negative_Infinity; j++){
				int _key = getKey(s, s, dir, 1, j);
				ParseForestItem tmp1 = chart[_key];
				chart[_key] = tmp;
				tmp = tmp1;
			}
			added = true;
			break;
		}
	}
	return added;
}

bool KBestParseForest::add(int s, int r, int t, int type, int dir, int comp, double score, ParseForestItem *p1, ParseForestItem *p2){
	bool added = false;
	int key = getKey(s, t, dir, comp, K-1);
	if(chart[key].prob > score){
		return false;
	}
	for(int i = 0;i< K; i++){
		key = getKey(s, t, dir, comp, i);
		if(chart[key].prob < score){
			ParseForestItem tmp = chart[key];
			chart[key] = ParseForestItem(s, r, t, type, dir, comp, score, p1, p2);
			for(int j = i + 1; j < K && tmp.prob > Negative_Infinity; j++){
				int _key = getKey(s, t, dir, comp, j);
				ParseForestItem tmp1 = chart[_key];
				chart[_key] = tmp;
				tmp = tmp1;
			}
			added = true;
			break;
		}
	}
	return added;
}

double KBestParseForest::getProb(int s, int t, int dir, int comp, int i){
	int key = getKey(s, t, dir, comp, i);
	return chart[key].prob;
}

double KBestParseForest::getProb(int s, int t, int dir, int comp){
	return getProb(s, t, dir, comp, 0);
}

vector<double>* KBestParseForest::getProbs(int s, int t, int dir, int comp){
	vector<double>* result = new vector<double>(K);
	for(int i = 0; i<K; i++){
		int key = getKey(s, t, dir, comp, i);
		(*result)[i] = chart[key].prob;
	}
	return result;
}

ParseForestItem KBestParseForest::getItem(int s, int t, int dir, int comp){
	return getItem(s, t, dir, comp, 0);
}

ParseForestItem KBestParseForest::getItem(int s, int t, int dir, int comp, int k){
	int key = getKey(s, t, dir, comp, k);
	return chart[key];
}

ParseForestItem* KBestParseForest::getItems(int s, int t, int dir, int comp){
	int key = getKey(s, t, dir, comp, 0);
	return (chart + key);
}

vector<pair<FeatureVector*, string*> >* KBestParseForest::getBestParses(DependencyInstance *inst, DependencyPipe *pipe){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	for(int i = 0; i < K; i++){
		int key = getKey(0, end, 0, 1, i);
		if(chart[key].prob > Negative_Infinity){
			(*d)[i].second = getDepString(chart[key]);
			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

vector<pair<FeatureVector*, string*> >* KBestParseForest::getBestParses(DependencyInstance *inst, DependencyPipe *pipe, vector<vector<vector<double> > > &probs, vector<vector<vector<vector<double> > > > &nt_probs, bool labeled){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	vector<ParseForestItem> result = vector<ParseForestItem>(K);
	for(int i = 0; i < K; i++){
		result[i] = ParseForestItem(0, -1, end, -1, 0, 1, Negative_Infinity, NULL, NULL);
	}
	for(int i = 1; i < end + 1; i++){
		ParseForestItem *left = getItems(1, i, 1, 1);
		ParseForestItem *right = getItems(i, end, 0, 1);
		vector<BestPair>* pairs = getKBestPairs(left, right);
		int size = (int)(pairs->size());
		bool added = true;
		double prob_0i = probs[0][i][0];
		double nt_prob_s_01 = 0.0;
		double nt_prob_t_00 = 0.0;
		if(labeled){
			nt_prob_s_01 = nt_probs[0][rootType][0][1];
			nt_prob_t_00 = nt_probs[i][rootType][0][0];
			prob_0i += nt_prob_s_01 + nt_prob_t_00;
		}

		for(int k = 0; k < size; k++){
			if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
				break;
			}

			int k1 = (*pairs)[k].first;
			int k2 = (*pairs)[k].second;
			double prob = left[k1].prob + right[k2].prob + prob_0i;
			added = false;
			for(int j = 0; j < K; j++){
				if(prob > result[j].prob){
					ParseForestItem tmp = result[j];
					result[j] = ParseForestItem(0, i, end, -1, 0, 1, prob, &(left[k1]), &(right[k2]));
					for(int m = j + 1; m < K; m++){
						ParseForestItem tmp1 = result[m];
						result[m] = tmp;
						tmp = tmp1;
					}
					added = true;
					break;
				}
			}
			if(!added){
				break;
			}
		}
		delete(pairs);
	}
	for(int i = 0; i < K; i++){
		if(result[i].prob > Negative_Infinity){
			string* str1 = getDepString(*(result[i].left));
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);
			int ll = (int)((*d)[i].second->length());
			if(ll != 0){
				(*d)[i].second->append(" ");
			}
			(*d)[i].second->append("0|");
			str1 = Util::intToString(result[i].r);
			str1->append(":");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = Util::intToString(KBestParseForest::rootType);
			str1->append(" ");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = getDepString(*(result[i].right));
			string* str2 = Util::trim(str1);
			pool->push_back(str1);
			(*d)[i].second->append(*str2);
			pool->push_back(str2);
			str1 = (*d)[i].second;
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);

			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

FeatureVector* KBestParseForest::getFeatureVector(DependencyInstance *inst, DependencyPipe *pipe, string *depString){
	FeatureVector* fv = NULL;
	int len = inst->length();
	vector<int> heads = vector<int>(len);
	vector<string*> labs = vector<string*>(len);
	vector<int> *heads_tmp = inst->heads;
	vector<string*> *labs_tmp = inst->deprels;
	heads[0] = (*heads_tmp)[0];
	labs[0] = (*labs_tmp)[0];

	vector<string*>* spans = Util::split(depString, ' ');
	for(int i = 1; i < len; ++i){
		size_t pos1 = (*spans)[i - 1]->find_first_of('|');
		size_t pos2 = (*spans)[i - 1]->find_first_of(':');
		size_t length = (*spans)[i - 1]->length();
		string tmp = (*spans)[i - 1]->substr(0, pos1);
		heads[i] = Util::stringToInt(&tmp);
		tmp = (*spans)[i - 1]->substr(pos2 + 1, length - pos2 - 1);
		int index = Util::stringToInt(&tmp);
		labs[i] = pipe->types[index];
	}
	inst->heads = &heads;
	inst->deprels = &labs;
	fv = pipe->createFeatureVector(inst, true);
	inst->heads = heads_tmp;
	inst->deprels = labs_tmp;
	int ll = (int)(spans->size());
	for(int i = 0; i < ll; i++){
		pool->push_back((*spans)[i]);
	}
	delete(spans);
	return fv;
}

string* KBestParseForest::getDepString(ParseForestItem &pfi){
	if(pfi.left == NULL){
		return pool->get();
	}

	if(pfi.comp == 1){
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		return result;
	}
	else if(pfi.dir == 0){
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		result->append(" ");
		str1 = Util::intToString(pfi.s);
		result->append(*str1 + "|");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.t);
		result->append(*str1 + ":");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else{
		string* str1 = Util::intToString(pfi.t);
		str1->append("|");
		string* str2 = Util::intToString(pfi.s);
		str2->append(":");
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::intToString(pfi.type);
		str2->append(" ");
		str1->append(*str2);
		pool->push_back(str2);
		string* result = str1;
		str1 = getDepString(*(pfi.left));
		str1->append(" ");
		str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::trim(str1);
		pool->push_back(str1);
		result->append(*str2);
		pool->push_back(str2);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
}

vector<BestPair>* KBestParseForest::getKBestPairs(ParseForestItem *item1, ParseForestItem *item2){
	bool* beenPushed = new bool[K * K];
	memset(beenPushed, 0, sizeof(bool) * K * K);

	vector<BestPair>* result = new vector<BestPair>(K, BestPair(-1, -1));

	BinaryHeap heap = BinaryHeap(K + 1);
	int n = 0;
	ValueIndexPair vip(item1[0].prob + item2[0].prob, 0, 0);
	heap.add(vip);
	beenPushed[0] = true;
	while(n < K){
		vip = heap.removeMax();
		if(vip.val <= Negative_Infinity){
			break;
		}
		
		(*result)[n].first = vip.i1;
		(*result)[n].second = vip.i2;
		n++;
		if(n >= K){
			break;
		}

		if(!beenPushed[(vip.i1 + 1) * K + vip.i2]){
			ValueIndexPair vp = ValueIndexPair(item1[vip.i1 + 1].prob +
				item2[vip.i2].prob, vip.i1 + 1, vip.i2);
			heap.add(vp);
			beenPushed[(vip.i1 + 1) * K + vip.i2] = true;
		}
		if(!beenPushed[vip.i1 * K + (vip.i2 + 1)]){
			ValueIndexPair vp = ValueIndexPair(item1[vip.i1].prob +
				item2[vip.i2 + 1].prob, vip.i1, vip.i2 + 1);
			heap.add(vp);
			beenPushed[vip.i1 * K + (vip.i2 + 1)] = true;
		}
	}
	delete[](beenPushed);
	return result;
}

ValueIndexPair::ValueIndexPair(double value, int i1, int i2){
	this->val = value;
	this->i1 = i1;
	this->i2 = i2;
}

ValueIndexPair::ValueIndexPair(){}

int ValueIndexPair::compareTo(ValueIndexPair &other){
	return (val < other.val) ? -1 : int(val > other.val);
}

BinaryHeap::BinaryHeap(int def_cap){
	this->DEFAULT_CAPACITY = def_cap;
	theArray = vector<ValueIndexPair>(DEFAULT_CAPACITY + 1);
	theArray[0] = ValueIndexPair(Positive_Infinity, -1, -1);
	currentSize = 0;
}

ValueIndexPair BinaryHeap::getMax(){
	return theArray[1];
}

int BinaryHeap::parent(int i){
	return i/2;
}

int BinaryHeap::leftChild(int i){
	return 2 * i;
}

int BinaryHeap::rightChild(int i){
	return 2 * i + 1;
}

void BinaryHeap::add(ValueIndexPair &e){
	int w = currentSize + 1;
	while(e.compareTo(theArray[parent(w)]) > 0){
		theArray[w] = theArray[parent(w)];
		w = parent(w);
	}
	theArray[w] = e;
	currentSize++;
}

ValueIndexPair BinaryHeap::removeMax(){
	ValueIndexPair min = theArray[1];
	theArray[1] = theArray[currentSize];
	currentSize--;
	bool switched = true;
	for(int p = 1; switched && p < currentSize;){
		switched = false;
		int leftC = leftChild(p);
		int rightC = rightChild(p);

		if(leftC <= currentSize){
			int largerC = leftC;
			if((rightC <= currentSize) && (theArray[rightC].compareTo(theArray[leftC]) > 0)){
				largerC = rightC;
			}
			if(theArray[largerC].compareTo(theArray[p]) > 0){
				ValueIndexPair temp = theArray[largerC];
				theArray[largerC] = theArray[p];
				theArray[p] = temp;
				p = largerC;
				switched = true;
			}
		}
	}

	return min;
}

KBestParseForest::~KBestParseForest(){
	delete[](chart);
}

